
import numpy as np
from quaternion import Quat, random_quaternion

class SimplexSO3():

    def __init__(self, func, xini=[], params=[]):

        self.func_call = func
        self.params = params

        self.reset_best()

        if xini == []:
            self.sim = np.r_[[Quat(0.,0,0).q,
                           Quat(.1,0,0).q,
                           Quat(0,.1,0).q,
                           Quat(0,0,.1).q]]
        else:
            self.reset_state(xini.q)

        self.fsim = np.array([self.func(x) for x in self.sim])
        self.sort()

    def reset_state(self,xiniq):
        xini = Quat(xiniq)
        self.sim = np.r_[[(Quat(-.003,-.003,-.003)*xini).q,
                           (Quat(.01,0,0)*xini).q,
                           (Quat(0,.01,0)*xini).q,
                           (Quat(0,0,.01)*xini).q]]

    def reset_best(self):
        self.best = np.array([0,0,0,0.])
        self.fbest = np.inf

    def func(self,x):
        v = self.func_call(x, *self.params)

        ## Update the register of the best value ever seen
        if v < self.fbest:
            self.best = x[:]
            self.fbest = v
        return v


    def sort(self):
        idx = np.argsort(self.fsim)
        self.sim = self.sim[idx]
        self.fsim = self.fsim[idx]

    def replace_worst(self,v,f,n=3):
        self.sim[n] = v.normalize().q
        self.fsim[n] = f
        self.sort()

    def errs(self):
        return np.max(np.abs(self.sim[1:]-self.sim[0])), self.fsim[3]-self.fsim[0]

    def opt_loop(self, iterations, xtol=1e-16, ftol=1e-16):
        ## Start the loop
        for it in range(iterations):
            ## Perform one optimization step
            self.opt_step()
            
            ## Update the best-seen register
            if self.fbest < self.fsim[0]:
                self.best = self.sim[0]
                self.fbest = self.fsim[0]

            ## Calculate errors in simplex shape and function values,
            ## break loop if we fall below the two tresholds
            xerr, ferr = self.errs()
            if xerr < xtol or ferr < ftol:
                break
        return it

    def opt_loop_sa(self, iterations, temperature, xtol=1e-16, ftol=1e-16):
        ## Start the loop
        for it in range(iterations):
            ## Perform one optimization step
            self.opt_step_sa(temperature)
            
            ## Update the best-seen register
            if self.fbest < self.fsim[0]:
                self.best = self.sim[0]
                self.fbest = self.fsim[0]

            ## Calculate errors in simplex shape and function values,
            ## break loop if we fall below the two tresholds
            xerr, ferr = self.errs()
            if xerr < xtol or ferr < ftol:
                break
        self.sort()

        return it

    def solution(self):
        ## Return the smallest vertex from the simplex. But you should consider to take self.best alternatively.
        return self.sim[0]

    def opt_step(self):
        debug = False
        def debug_print(x):
            if debug:
                print x
        # print 70*'-'
        # print self.sim
        # print self.fsim
        ## Calculate centroid
        #cen = Quat(self.sim[0]+self.sim[1]+self.sim[2]).normalize()
        cen = Quat(self.sim[0]+self.sim[1]).normalize()

        ## REFLECT
        ## Calculate reflection and extension vertices and function values
        delta = cen/Quat(self.sim[3])
        vr = delta * cen
        fr = self.func(vr.q)

        debug_print(vr)
        debug_print(fr)

        if fr >= self.fsim[0] and fr < self.fsim[2]:
            ## Reflection result is good, but not best. We go with that then.
            debug_print('R')
            self.replace_worst(vr,fr)
            return

        ## EXPANSION
        ve = delta * vr  ## delta * delta * cen
        fe = self.func(ve.q)

        if fr < self.fsim[0]:
            if fe < fr:
                ## If the expansion turns out to be even better, use that.
                debug_print('E')
                self.replace_worst(ve,fe)
                return
            else:
                ## Reflection was still a pretty good move, so go with that.
                debug_print('ER')
                self.replace_worst(vr,fr)
                return

        ## At this point reflection was a bad move, making the worst
        ## value as bad as the second worst, at best.

        if fr < self.fsim[3]:
            ## CONTRACTION
            vc = delta.sqrt() * cen
            fc = self.func(vc.q)

            if fc <= fr:
                debug_print('C')
                self.replace_worst(vc,fc)
                return
        else: # fr >= self.fsim[3]:
            ## INTERNAL CONTRACTION
            vcc = cen / delta.sqrt()
            fcc = self.func(vcc.q)

            if fcc <= fr:
                debug_print('CC')
                self.replace_worst(vcc,fcc)
                return            

        ## REDUCTION
        debug_print('D')
        for k in range(1,4):
            self.sim[k] = ((Quat(self.sim[k])/Quat(self.sim[0])).sqrt() * Quat(self.sim[0])).q
            self.fsim[k] = self.func(self.sim[k])
        self.sort()

    def opt_step_sa(self, temperature = 0.0):
        debug = False
        # print 70*'-'
        # print self.sim
        # print self.fsim

        ## Function values with random noise for this iteration
        fsim = self.fsim + temperature * -np.log(np.random.rand(4))
        ## The vertex order for these values
        idx = np.argsort(fsim)
        
        ## Calculate centroid
        cen = Quat(self.sim[idx[0]]+self.sim[idx[1]]+self.sim[idx[2]]).normalize()

        ## REFLECT
        ## Calculate reflection and extension vertices and function values
        delta = cen/Quat(self.sim[idx[3]])
        vr = delta * cen
        fr = self.func(vr.q)
        hot_fr = fr - temperature * -np.log(np.random.rand())

        if hot_fr >= fsim[0] and hot_fr < fsim[2]:
            ## Reflection result is good, but not best. We go with that then.
            if debug:
                print 'R'
            self.replace_worst(vr,fr,idx[3])
            return

        ## EXPANSION
        ve = delta * vr  ## delta * delta * cen
        fe = self.func(ve.q)
        hot_fe = fe - temperature * -np.log(np.random.rand())

        if hot_fr < fsim[0]:
            ## If the reflection is amazingly good, try to keep moving.
            if hot_fe < hot_fr:
                ## If the expansion turns out to be even better, use that.
                if debug:
                    print 'E'
                self.replace_worst(ve,fe,idx[3])
                return
            else:
                ## Reflection was still a pretty good move, so go with that.
                if debug:
                    print 'ER'
                self.replace_worst(vr,fr)
                return

        ## At this point reflection was a bad move, making the worst
        ## value as bad as the second worst, at best.

        ## CONTRACTION
        vc = delta.sqrt() * Quat(self.sim[idx[3]])  ## i.e., cen / delta.sqrt()
        fc = self.func(vc.q)
        hot_fc = fc - temperature * -np.log(np.random.rand())

        if hot_fc < fsim[3]:
            if debug:
                print 'C'
            self.replace_worst(vc,fc,idx[3])
            return

        ## REDUCTION
        if debug:
            print 'D'
        for k in range(1,4):
            self.sim[idx[k]] = ((Quat(self.sim[idx[k]])/Quat(self.sim[idx[0]])).sqrt() * Quat(self.sim[idx[0]])).q
            self.fsim[k] = self.func(self.sim[idx[k]])
        self.sort()




if __name__ == '__main__':
    from pylab import *

    def func(q):
        return -np.abs(dot(array([1,2,3,4]),q)).sum()

    # def func(q):
    #     r = Quat(q).rot().T
    #     e1 = np.abs(np.dot(np.array([1.,0,0]), r[2]))
    #     e2 = np.abs(np.dot(np.array([0.,1,0]), r[0]))
    #     e3 = np.abs(np.dot(np.array([.707,.707,0]), r[2]))
    #     return e1+e2+e3


    ss = SimplexSO3(func, Quat(0,0,0))

    temp = 1
    var = 0.05
    m = array([0.,0.,0.])
    k = array([.999, .99, .9])
    kt = .9999
    while True:
        ss.opt_step_sa(temp)
        ex, ef = ss.errs()
        m = k*m + (1.0-k) * ex
        temp = temp*kt + (1-kt) * 100*(var - m[1])
        if temp < 0:
            temp = 0

        print temp, var*100, m*100, ex


